"""
测试二：沪深大盘实时数据服务
从网上获取沪深大盘实时价格数据，每隔3秒获取一次，保存到Redis
"""

import asyncio
import json
import logging
import random
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Set
import redis.asyncio as redis
from fastapi import WebSocket
import akshare as ak

logger = logging.getLogger(__name__)

class MarketIndexService:
    def __init__(self):
        self.redis_client: Optional[redis.Redis] = None
        self.active_connections: Set[WebSocket] = set()
        self.subscribed_stocks: Set[str] = set()
        self.is_running = False
        
    async def init_redis(self):
        """初始化Redis连接"""
        try:
            self.redis_client = redis.Redis(
                host='localhost',
                port=6380,
                decode_responses=True,
                socket_connect_timeout=5,
                socket_timeout=5
            )
            await self.redis_client.ping()
            logger.info("✅ Redis连接成功")
        except Exception as e:
            logger.warning(f"⚠️ Redis连接失败，使用内存缓存: {e}")
            self.redis_client = None
    
    async def connect_websocket(self, websocket: WebSocket):
        """添加WebSocket连接"""
        await websocket.accept()
        self.active_connections.add(websocket)
        logger.info(f"📱 新的WebSocket连接，当前连接数: {len(self.active_connections)}")
        
    async def disconnect_websocket(self, websocket: WebSocket):
        """移除WebSocket连接"""
        self.active_connections.discard(websocket)
        logger.info(f"📱 WebSocket连接断开，当前连接数: {len(self.active_connections)}")
    
    async def subscribe_stock(self, stock_code: str):
        """订阅股票实时数据"""
        self.subscribed_stocks.add(stock_code)
        logger.info(f"📊 订阅股票: {stock_code}，当前订阅数: {len(self.subscribed_stocks)}")
    
    async def unsubscribe_stock(self, stock_code: str):
        """取消订阅股票"""
        self.subscribed_stocks.discard(stock_code)
        logger.info(f"📊 取消订阅股票: {stock_code}")
    
    async def get_realtime_price(self, stock_code: str) -> Dict:
        """获取实时股价数据"""
        try:
            # 首先尝试从Redis缓存获取
            if self.redis_client:
                cached_data = await self.redis_client.get(f"realtime:{stock_code}")
                if cached_data:
                    data = json.loads(cached_data)
                    # 检查缓存是否过期（30秒）
                    cache_time = datetime.fromisoformat(data['timestamp'])
                    if datetime.now() - cache_time < timedelta(seconds=30):
                        return data
            
            # 获取实时数据（演示模式：生成模拟数据）
            price_data = await self._fetch_mock_realtime_data(stock_code)
            
            # 缓存到Redis
            if self.redis_client:
                await self.redis_client.setex(
                    f"realtime:{stock_code}", 
                    30,  # 30秒过期
                    json.dumps(price_data, ensure_ascii=False)
                )
            
            return price_data
            
        except Exception as e:
            logger.error(f"获取实时数据失败 {stock_code}: {e}")
            return await self._fetch_mock_realtime_data(stock_code)
    
    async def _fetch_mock_realtime_data(self, stock_code: str) -> Dict:
        """生成模拟实时数据（演示用）"""
        # 基础价格（根据股票代码生成）
        base_price = hash(stock_code) % 100 + 10
        
        # 模拟价格波动
        change_percent = random.uniform(-3, 3)
        current_price = base_price * (1 + change_percent / 100)
        change_amount = current_price - base_price
        
        # 模拟成交量
        volume = random.randint(1000000, 50000000)
        
        return {
            "stock_code": stock_code,
            "stock_name": await self._get_stock_name(stock_code),
            "current_price": round(current_price, 2),
            "change_amount": round(change_amount, 2),
            "change_percent": round(change_percent, 2),
            "volume": volume,
            "timestamp": datetime.now().isoformat(),
            "high": round(current_price * 1.02, 2),
            "low": round(current_price * 0.98, 2),
            "open": round(base_price, 2),
            "prev_close": round(base_price, 2)
        }
    
    async def _get_stock_name(self, stock_code: str) -> str:
        """获取股票名称"""
        # 扩展的股票名称映射表
        stock_names = {
            # 主板股票
            "000001": "平安银行",
            "000002": "万科A",
            "000858": "五粮液",
            "000876": "新希望",

            # 深圳主板
            "002962": "五方光电",
            "002415": "海康威视",
            "002594": "比亚迪",
            "002230": "科大讯飞",

            # 上海主板
            "600519": "贵州茅台",
            "600036": "招商银行",
            "600000": "浦发银行",
            "601318": "中国平安",
            "601166": "兴业银行",

            # 创业板
            "300750": "宁德时代",
            "300920": "润阳科技",
            "300059": "东方财富",
            "300142": "沃森生物",

            # 科创板
            "688981": "中芯国际",
            "688599": "天合光能",
            "688111": "金山办公",
            "688036": "传音控股"
        }

        # 如果在映射表中找到，返回对应名称
        if stock_code in stock_names:
            return stock_names[stock_code]

        # 根据股票代码规律推测名称
        if stock_code.startswith("000"):
            return f"深圳主板{stock_code}"
        elif stock_code.startswith("002"):
            return f"深圳中小板{stock_code}"
        elif stock_code.startswith("300"):
            return f"创业板{stock_code}"
        elif stock_code.startswith("600") or stock_code.startswith("601") or stock_code.startswith("603"):
            return f"上海主板{stock_code}"
        elif stock_code.startswith("688"):
            return f"科创板{stock_code}"
        else:
            return f"股票{stock_code}"
    
    async def broadcast_to_websockets(self, data: Dict):
        """向所有WebSocket连接广播数据"""
        if not self.active_connections:
            return
            
        message = json.dumps(data, ensure_ascii=False)
        disconnected = set()
        
        for websocket in self.active_connections:
            try:
                await websocket.send_text(message)
            except Exception as e:
                logger.warning(f"WebSocket发送失败: {e}")
                disconnected.add(websocket)
        
        # 清理断开的连接
        for websocket in disconnected:
            self.active_connections.discard(websocket)
    
    async def start_realtime_updates(self):
        """启动实时数据更新任务"""
        if self.is_running:
            return
            
        self.is_running = True
        logger.info("🚀 启动实时股价数据推送服务")
        
        while self.is_running:
            try:
                if self.subscribed_stocks and self.active_connections:
                    # 获取所有订阅股票的实时数据
                    tasks = [
                        self.get_realtime_price(stock_code) 
                        for stock_code in self.subscribed_stocks
                    ]
                    
                    if tasks:
                        results = await asyncio.gather(*tasks, return_exceptions=True)
                        
                        # 广播数据
                        for result in results:
                            if isinstance(result, dict):
                                await self.broadcast_to_websockets({
                                    "type": "realtime_price",
                                    "data": result
                                })
                
                # 每3秒更新一次
                await asyncio.sleep(3)
                
            except Exception as e:
                logger.error(f"实时数据更新错误: {e}")
                await asyncio.sleep(5)
    
    async def stop_realtime_updates(self):
        """停止实时数据更新"""
        self.is_running = False
        logger.info("⏹️ 停止实时股价数据推送服务")
    
    async def get_stock_history_minute(self, stock_code: str, period: str = "1d") -> List[Dict]:
        """获取分钟级历史数据"""
        try:
            # 演示模式：生成模拟分钟数据
            data = []
            base_price = hash(stock_code) % 100 + 10
            current_time = datetime.now()
            
            # 生成最近60分钟的数据
            for i in range(60, 0, -1):
                timestamp = current_time - timedelta(minutes=i)
                price_change = random.uniform(-2, 2)
                price = base_price * (1 + price_change / 100)
                
                data.append({
                    "timestamp": timestamp.isoformat(),
                    "price": round(price, 2),
                    "volume": random.randint(100000, 1000000)
                })
            
            return data
            
        except Exception as e:
            logger.error(f"获取分钟数据失败 {stock_code}: {e}")
            return []

# 全局实例
realtime_service = MarketIndexService()
